import cv2
video_path = './vtest.avi'

#角点检测
feature_params = dict(maxCorners=100,
                      qualityLevel=0.3,
                      minDistance=7,
                      blockSize=7)

#l-k
lk_params = dict(winSize=(15, 15),
                 maxLevel=2,
                 criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10,
                           0.03))
cap = cv2.VideoCapture(video_path)

#第一帧特征值
ret, prev = cap.read()
img_gray = cv2.cvtColor(prev, cv2.COLOR_BGR2GRAY)

p0 = cv2.goodFeaturesToTrack(img_gray, mask=None, **feature_params)

while True:
    ret, frame = cap.read()
    if not ret:
        break
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # 计算光流
    p1, st, err = cv2.calcOpticalFlowPyrLK(img_gray, gray, p0, None,
                                           **lk_params)

    #选取好的跟踪点
    goodPoints = p1[st == 1]
    goodPrevPoints = p0[st == 1]

    res = frame.copy()
    drawColor = (0, 0xff, 0)
    for i, (cur, prev) in enumerate(zip(goodPoints, goodPrevPoints)):
        x0, y0 = cur.ravel()
        x1, y1 = prev.ravel()
        cv2.line(res, (x0, y0), (x1, y1), drawColor)
        cv2.circle(res, (x0, y0), 3, drawColor)
    #更新
    img_grey = gray.copy()
    p0 = goodPoints.reshape(-1, 1, 2)

    cv2.imshow('Result', res)

    key = cv2.waitKey(30)
    if key == 27:
        break
cap.release()
cv2.destroyAllWindows()
